#include "component_clustering.h"
#include <queue>
#include <cmath>
#include <cstring>
#include <set>

using namespace std;
using namespace pcl;

// --- 工具函数：点投影到网格索引 ---
inline bool pointToGrid(float x, float y, int& xi, int& yi) {
    float xC = x + roiM / 2;
    float yC = y + roiM / 2;
    if (xC < 0 || xC >= roiM || yC < 0 || yC >= roiM)
        return false; // 超出ROI
    xi = std::floor(numGrid * xC / roiM);
    yi = std::floor(numGrid * yC / roiM);
    if(xi<0 || xi>=numGrid || yi<0 || yi>=numGrid) return false;
    return true;
}

// --- 主聚类函数 ---
void componentClustering(
    const pcl::PointCloud<pcl::PointXYZI>::Ptr& cloud,
    std::array<std::array<int, numGrid>, numGrid>& cartesianData,
    int& numCluster)
{
    // 1. 网格数据初始化
    cartesianData.fill({});
    for(auto& row:cartesianData) row.fill(0);

    // 2. 投影点云到网格，统计每格点数
    std::array<std::array<int, numGrid>, numGrid> occupancy;
    occupancy.fill({});
    for(auto& row: occupancy) row.fill(0);

    for(const auto& pt : cloud->points) {
        int xi, yi;
        if(pointToGrid(pt.x, pt.y, xi, yi)){
            occupancy[xi][yi]++;
        }
    }

    // 3. 选定阈值，将有效格子置为-1（候选障碍物）
    const int minPointsInCell = 2; // 单格至少点数
    for(int i=0;i<numGrid;++i) {
        for(int j=0;j<numGrid;++j) {
            if(occupancy[i][j]>=minPointsInCell)
                cartesianData[i][j] = -1; // 有障碍
            else
                cartesianData[i][j] = 0; // 背景
        }
    }

    // 4. 连通域Flood fill聚类，每块障碍物分配独立ID
    int current_cluster = 0;
    int dx[8] = {-1,-1,-1,0,0,1,1,1};
    int dy[8] = {-1,0,1,-1,1,-1,0,1};
    for(int i=0;i<numGrid;++i) {
        for(int j=0;j<numGrid;++j) {
            if(cartesianData[i][j]==-1) {
                ++current_cluster;
                queue<pair<int,int>> q;
                q.push({i,j});
                cartesianData[i][j]=current_cluster;
                while(!q.empty()) {
                    auto [x,y] = q.front(); q.pop();
                    for(int d=0; d<8; ++d){
                        int nx = x+dx[d];
                        int ny = y+dy[d];
                        if(nx>=0 && nx<numGrid && ny>=0 && ny<numGrid){
                            if(cartesianData[nx][ny]==-1){
                                cartesianData[nx][ny]=current_cluster;
                                q.push({nx,ny});
                            }
                        }
                    }
                }
            }
        }
    }
    numCluster = current_cluster;
}

// --- 拆分为单独点云 ---
void getClusteredPoints(
    const pcl::PointCloud<pcl::PointXYZI>::Ptr& cloud,
    const std::array<std::array<int, numGrid>, numGrid>& cartesianData,
    std::vector<pcl::PointCloud<pcl::PointXYZI>>& clusteredPoints
){
    // 首先统计聚类个数
    std::set<int> cluster_ids;
    for(int i=0;i<numGrid;++i)
        for(int j=0;j<numGrid;++j)
            if(cartesianData[i][j]>0) cluster_ids.insert(cartesianData[i][j]);

    int n = cluster_ids.size();
    if(n==0) return;
    clusteredPoints.clear();
    clusteredPoints.resize(*cluster_ids.rbegin());

    for(const auto& pt : cloud->points){
        int xi, yi;
        if(pointToGrid(pt.x, pt.y, xi, yi)){
            int cid = cartesianData[xi][yi];
            if(cid>0){
                clusteredPoints[cid-1].push_back(pt); // id从1开始
            }
        }
    }
}